`include "defines.v"

module EX_Stage (
    input  wire [`BUSLEN-1:0]           pc_i,
    input  wire                         pc_valid_i,
    output wire [`BUSLEN-1:0]           pc_o,
    output wire                         pc_valid_o,

    // ALU Oprends
    input  wire [`BUSLEN-1:0]           regbus_A,
    input  wire [`BUSLEN-1:0]           regbus_B,
    input  wire [`BUSLEN-1:0]           imm,
    output wire [`BUSLEN-1:0]           ALU_out,
    output wire [ 3: 0]                 psw_flags,

    // ALU Operation Control
    input  wire [ 1: 0]                 ALU_oprend_A_src,
    input  wire [ 1: 0]                 ALU_oprend_B_src,
    input  wire [ 2: 0]                 ALU_op,
    input  wire [ 2: 0]                 ALU_out_ext_type,

    // Memory Access
    input  wire                         mem_r_en_i,
    input  wire [ 2: 0]                 mem_r_ext_type_i,
    input  wire                         mem_w_en_i,
    input  wire [ 1: 0]                 mem_w_size_i,

    output wire                         mem_r_en_o,
    output wire [ 2: 0]                 mem_r_ext_type_o,
    output wire                         mem_w_en_o,
    output wire [ 1: 0]                 mem_w_size_o,
    output wire [`BUSLEN-1:0]           regbus_B_o,

    // Regfile/CSR R/W
    input  wire [`REGFILE_ADDR_LEN-1:0] rd_w_addr_i,
    input  wire                         rd_w_en_i,
    input  wire                         rd_w_src_i,

    output wire [`REGFILE_ADDR_LEN-1:0] rd_w_addr_o,
    output wire                         rd_w_en_o,
    output wire                         rd_w_src_o,

    input  wire                         csr_w_en_i,
    input  wire [`CSR_ADDR_LEN-1:0]     csr_w_addr_i,

    output wire                         csr_w_en_o,
    output wire [`CSR_ADDR_LEN-1:0]     csr_w_addr_o,

    // Jump/Branch Signals
    input  wire                         adder_src,
    input  wire                         inst_jump_i,
    input  wire [ 2: 0]                 inst_branch_type_i,

    output wire                         inst_jump_o,
    output wire [ 2: 0]                 inst_branch_type_o,

    // Privileged Instruction / Exception
    input  wire                         is_system_inst_i,
    output wire                         is_system_inst_o,

    input  wire                         inst_ecall_i,
    input  wire                         inst_mret_i,

    output wire                         inst_ecall_o,
    output wire                         inst_mret_o
);

    //----------ALU----------//
    // Oprends Generate
    reg  [`XLEN-1:0] ALU_oprend_A;
    reg  [`XLEN-1:0] ALU_oprend_B;
    reg  [`XLEN:0]   ALU_out_ext;

    always @(*) begin
        case (ALU_oprend_A_src)
            `PC                   : ALU_oprend_A = pc_i;
            `REGBUS_A             : ALU_oprend_A = regbus_A;
            `ZEXT_BITSEL_REGBUS_A : ALU_oprend_A = { 32'b0 , regbus_A[31:0] };
            `SEXT_BITSEL_REGBUS_A : ALU_oprend_A = { {32{regbus_A[31]}} , regbus_A[31:0] };
            default               : ALU_oprend_A = regbus_A;
        endcase
    end

    always @(*) begin
        case (ALU_oprend_B_src)
            `REGBUS_B             : ALU_oprend_B = regbus_B;
            `IMM                  : ALU_oprend_B = imm;
            `CONST_4              : ALU_oprend_B = 64'b0100;
            `ZEXT_BITSEL_REGBUS_B : ALU_oprend_B = ALU_out_ext_type[0] ?  { 59'b0 , regbus_B[4:0]}  :  { 58'b0 , regbus_B[5:0] }  ;
            default               : ALU_oprend_B = regbus_B;
        endcase
    end

    // Operation Executation
    always @(*) begin
        case (ALU_op)
            `ADD_OP : ALU_out_ext = { 1'b0 , ALU_oprend_A } + { 1'b0 , ALU_oprend_B };
            `SUB_OP : ALU_out_ext = { 1'b0 , ALU_oprend_A } - { 1'b0 , ALU_oprend_B };
            `AND_OP : ALU_out_ext = { 1'b0 , ALU_oprend_A } & { 1'b0 , ALU_oprend_B };
            `OR_OP  : ALU_out_ext = { 1'b0 , ALU_oprend_A } | { 1'b0 , ALU_oprend_B };
            `XOR_OP : ALU_out_ext = { 1'b0 , ALU_oprend_A } ^ { 1'b0 , ALU_oprend_B };
            `SLL_OP : ALU_out_ext = { 1'b0 , ALU_oprend_A } << { 1'b0 , ALU_oprend_B };
            `SRL_OP : ALU_out_ext = { 1'b0 , ALU_oprend_A } >> { 1'b0 , ALU_oprend_B };
            `SRA_OP : ALU_out_ext = $signed({ ALU_oprend_A[63] , ALU_oprend_A }) >>> { 1'b0 , ALU_oprend_B };
            default : ALU_out_ext = { 1'b0 , ALU_oprend_A } + { 1'b0 , ALU_oprend_B };
        endcase
    end

    assign ALU_out = ALU_out_ext_type[2] ?  (ALU_out_ext_type[1] ?  { 63'b0 , `CF }  :  { 63'b0 , `SF^`OF })     :
                     ALU_out_ext_type[0] ?  { {32{ALU_out_ext[31]}} , ALU_out_ext[31:0] }  :  ALU_out_ext[63:0]  ;

    // PSW Flags
    assign `CF = ALU_out_ext[64];
    assign `SF = ALU_out_ext[63];
    assign `OF = (ALU_oprend_A[63] & ~ALU_oprend_B[63] & ~ALU_out_ext[63])|(~ALU_oprend_A[63] & ALU_oprend_B[63] & ALU_out_ext[63]);
    assign `ZF = ~(|ALU_out_ext);

    //----------Adder----------//
    wire [`BUSLEN-1:0] adder_oprendA;
    wire [`BUSLEN-1:0] adder_oprendB;

    assign adder_oprendA = adder_src ?  regbus_A  :  pc_i  ;
    assign adder_oprendB = imm;
    assign pc_o = adder_oprendA + adder_oprendB;

    //----------Branch/Jump Signals----------//
    assign pc_valid_o         = pc_valid_i;
    assign inst_jump_o        = inst_jump_i;
    assign inst_branch_type_o = inst_branch_type_i;

    //----------Privileged Instructions----------//
    assign is_system_inst_o = is_system_inst_i;
    assign inst_ecall_o     = inst_ecall_i;
    assign inst_mret_o      = inst_mret_i;

    //----------Memory/Regfile/CSR Access----------//
    // Memory
    assign mem_r_en_o       = mem_r_en_i;
    assign mem_r_ext_type_o = mem_r_ext_type_i;
    assign mem_w_en_o       = mem_w_en_i;
    assign mem_w_size_o     = mem_w_size_i;
    assign regbus_B_o       = regbus_B;
    
    // Regfile
    assign rd_w_en_o   = rd_w_en_i;
    assign rd_w_src_o  = rd_w_src_i;
    assign rd_w_addr_o = rd_w_addr_i;

    // CSR
    assign csr_w_en_o   = csr_w_en_i;
    assign csr_w_addr_o = csr_w_addr_i;
    
endmodule